/*
 * Filter.cpp
 *
 *  Created on: Apr 23, 2011
 *      Author: tflanzer
 */

#include "Filter.h"
#include <iostream>
using namespace std;

const int ione = 1;
const double done = 0;


// Initialization
Filter::Filter(Aircraft* aci) {

//	for (int i = 0; i < numStates; i++) {
//		for (int j = 0; j < numStates;j++) {
//			if (i == j) {
//				P[i][j] = 1;
//			} else if (j > i){
//				P[i][j] = (double)(i+1)/(double)(j+1);
//			} else if (i > j) {
//				P[i][j] = (double)(j+1)/(double)(i+1);
//			}
//		}
//	}
	for (int i = 0; i<numStates; i++) {
		for (int j = 0; j<numStates; j++) {
			if (i == j) {
				P[i][j] = 1.0;
			} else {
				P[i][j] = 0.0;
			}
		}
	}
	
	Q_IMU_GPS = {10.0,10.0,25.0,0.1,0.1,0.04,0.04,0.04,0.0001,0.0001,0.0001,100.0,100.0};
	Q_IMU = {0.04,0.04,0.04,0.0001,0.0001,0.0001,100.0,100.0};
	// N,E,D,u,v,w (inertial velocites in body frame),
	//e0,e1,e2,e3,p,q,r,WindN,WindE,WindD,abx,aby,abz,wbx,wby,wbz,barob,diffb
	R = {1.0,1.0,1.0,0.01,0.01,0.01,0.0001,0.0001,0.0001,0.0001,
			0.01,0.01,0.01,1.0,1.0,1.0,0.001,0.001,0.001,0.001,0.001,0.001,0.01,0.01};
	
	// UKF Parameters
	alpha = 0.01; 
	beta = 2.0;
	kappa = 0.0;
	lambda = alpha*alpha*(numStates+kappa)-numStates;
	gamma = sqrt(numStates+lambda);

	// Sigma Point Weights
	for (int i = 0; i<2*numStates+1; i++) {
		wm[i] = 1.0/(2.0*(numStates+lambda));
		wc[i] = wm[i];
	}
	wm[0] = lambda/(numStates+lambda);
	wc[0] = wm[0] + (1-alpha*alpha + beta);

	ac = aci;
	dt = 0.01;
}

void Filter::AUPdataReceived(Communication &com) {
	
	// Determine packet type
	bool flag = true;
	if (com.imuin) {
		dt = com.tIMU[1] - com.tIMU[0];
		/*
		 *   Nothing happens when a GPS packet is received. We always wait for the
		 *   next IMU packet before performing an update. So check to see if the
		 *   latest GPS time is between the current IMU time and the previous IMU
		 *   time. If so, perform the combined measurement update (flag = false).
		 */
		if ((com.tIMU[1] > com.tGPS[1]) && (com.tIMU[0] < com.tGPS[1])) {
			flag = false;
		}

		// Write current state
		double *xcurr;
		xcurr = (double*)malloc(sizeof(double)*numStates);
		getCurrentState(xcurr);

		// 1. Generate sigma points
		double *sigmaPts;
		sigmaPts = (double*)malloc(sizeof(double)*numStates*(2*numStates+1));
		generateSigmaPointsProcess(sigmaPts,xcurr);

		// 2. Propagate points through process model
		double *xprocess;
		xprocess = (double*)malloc(sizeof(double)*numStates*(2*numStates+1));
		propagatePtsThroughProcessModel(xprocess,sigmaPts);

		// 3. Compute predicted mean
		double *mubt;
		mubt = (double*)malloc(sizeof(double)*numStates);
		computePredictedMean(xprocess,mubt);

		// 4. Compute predicted variance
		double *Sbart;
		Sbart = (double*)malloc(sizeof(double)*numStates*(2*numStates+1));
		computePredictedVariance(xprocess,mubt,Sbart);

		// 5. Calculate sigma points from prediction step
		generateSigmaPointsMeasurement(sigmaPts,xcurr,Sbart);

		// 6. Compute predicted observation for each sigma point
		double *xmeas;
		xmeas = (double*)malloc(sizeof(double)*numStates*(2*numStates+1));
		computePredictedObservation(xmeas, sigmaPts, flag);

		// 7. Compute predicted measurement
		double *zbt;
		if (flag) {
			double* zbt = (double*)malloc(sizeof(double)*imuSize);
			computePredictedMeasurement(xmeas,zbt,true);
		}
		else {
			zbt = (double*)malloc(sizeof(double)*(gpsSize+imuSize));
			computePredictedMeasurement(xmeas,zbt,false);
		}

		// 8. Compute predicted measurement uncertainty
		double *St;
		if (flag) {
			St = (double*)malloc(sizeof(double)*imuSize*imuSize);
			computePredictedMeasurementUncertainty(xmeas,zbt,St,true);
		} else {
			St = (double*)malloc(sizeof(double)*(imuSize+gpsSize)*(imuSize+gpsSize));
			computePredictedMeasurementUncertainty(xmeas,zbt,St,false);
		}

		// 9. Compute cross-covariance
		double *Scc;
		if (flag) {
			Scc = (double*)malloc(sizeof(double)*numStates*imuSize);
			computeCrossCovariance(xprocess, xmeas, mubt, zbt, Scc, true);
		} else {
			Scc = (double*)malloc(sizeof(double)*numStates*(imuSize+gpsSize));
			computeCrossCovariance(xprocess, xmeas, mubt, zbt, Scc, false);
		}

		// 10. Compute Kalman gain
		double *Kt;
		if (flag) {
			Kt = (double*)malloc(sizeof(double)*numStates*imuSize);
			computeKalmanGain(St, Scc, Kt, true);
		} else {
			Kt = (double*)malloc(sizeof(double)*numStates*(imuSize+gpsSize));
			computeKalmanGain(St, Scc, Kt, false);
		}

		// 11. Compute new mean and covariance
		computeNewMeanandCovariance(mubt,Kt,zbt,Sbart,St,flag);

		free(xcurr); free(sigmaPts); free(xprocess); free(Sbart); free(xmeas);
		free(zbt); free(St); free(Scc);
	}
}

// Destructor - do not change
Filter::~Filter() {
}

void Filter::computeNewMeanandCovariance(double *mubt, double *Kt, double *zbt,
		double *Sbart, double *St, bool flag) {

	int N, M;
	M = numStates;
	double dzero = 0.0;

	if (flag) { // IMU
		N = imuSize;
	} else { // GPS+IMU
		N = gpsSize + imuSize;
	}

	double *zt  = (double*)malloc(sizeof(double)*N);
	double *mut = (double*)malloc(sizeof(double)*M);
	double *tmp = (double*)malloc(sizeof(double)*M);
	double *Et  = (double*)malloc(sizeof(double)*M*M);
	double *Ktz = (double*)malloc(sizeof(double)*M);
	double *StKtt = (double*)malloc(sizeof(double)*M*N);

	if (flag) {
		*(zt+0) = apdat.ax;
		*(zt+1) = apdat.ay;
		*(zt+2) = apdat.az;
		*(zt+3) = apdat.p;
		*(zt+4) = apdat.q;
		*(zt+5) = apdat.r;
		*(zt+6) = apdat.baropress;
		*(zt+7) = apdat.diffpress;
	} else {
		*(zt+0) = apdat.ax;
		*(zt+1) = apdat.ay;
		*(zt+2) = apdat.az;
		*(zt+3) = apdat.p;
		*(zt+4) = apdat.q;
		*(zt+5) = apdat.r;
		*(zt+6) = apdat.baropress;
		*(zt+7) = apdat.diffpress;
		*(zt+8) = apdat.GPSlat;
		*(zt+9) = apdat.GPSlon;
		*(zt+10) = apdat.GPSgndspeed;
		*(zt+11) = apdat.GPScog;
		*(zt+12) = apdat.GPSalt;
	}

	for (int i = 0; i < numStates; i++) {
		*(tmp+i) = (*(zt+i)-*(zbt+i));
	}

	// Kt*(zt-zbt)
	dgemv_("N", &M, &N, &done, Kt, &M, tmp, &ione, &dzero, Ktz, &ione);

	// Compute new estimate
	for (int i = 0; i < numStates; i++) {
		*(mut+i) = *(mubt+i) + *(Ktz+i);
	}

	updateState(mut);

	// St*Kt^T
	dgemm_("N","T", &N, &N, &N, &done, St, &N, Kt, &M, &dzero, StKtt, &N);

	// Kt*(St*Kt^T)
	dgemm_("N","N", &N, &M, &N, &done, Kt, &M, StKtt, &N, &dzero, Et, &M);

	// Compute updated covariance
	for (int i = 0; i < M; i++) {
		for (int j = 0; j < M; j++) {
			P[i][j] = *(Sbart+j+i*M)-*(Et+i+j*M);
		}
	}

	free(mut); free(tmp); free(Et); free(Ktz); free(StKtt); free(zt);
}

void Filter::computeKalmanGain(double *St, double *Scc, double *Kt, bool flag) {

	int N, M, info;
	M = numStates;

	if (flag) { // IMU
		N = imuSize;
	} else { // GPS+IMU
		N = gpsSize + imuSize;
	}

	double *St_trans  = (double*)malloc(sizeof(double)*N*N);
	double *Scc_trans = (double*)malloc(sizeof(double)*M*N);
	int *ipiv = (int*)malloc(sizeof(int)*N);

	// Compute transpose
	for (int i = 0; i < N; i++) {
		for (int j = 0; j < N; j++) {
			*(St_trans+i+j*N) = *(St+j+i*N);
		}
	}
	// Compute transpose
	for (int i = 0; i < N; i++) {
		for (int j = 0; j < M; j++) {
			*(Scc_trans+i+j*N) = *(Scc+j+i*N);
		}
	}

	dgesv_(&N, &M, St_trans, &N, ipiv, Scc_trans, &N, &info);

	// Compute transpose
	for (int i = 0; i < N; i++) {
		for (int j = 0; j < M; j++) {
			*(Kt+j+i*N) = *(Scc_trans+i+j*N);
		}
	}

	free(ipiv); free(St_trans); free(Scc_trans);
}

void Filter::computeCrossCovariance(double *xprocess, double *xmeas, double *mubt,
		double *zbt, double *Scc, bool flag) {

	int N, M;
	double wca;
	M = numStates;

	if (flag) { // IMU
		N = imuSize;
	} else { // GPS+IMU
		N = gpsSize + imuSize;
	}

	double *vec1 = (double*)malloc(sizeof(double)*N);
	double *vec2 = (double*)malloc(sizeof(double)*numStates);

	for (int j = 0; j < numStates; j++) {
		for (int i = 0; i < N; i++) {
			*(Scc+i+j*N) = 0.0;
		}
	}

	for (int i = 0; i < 2*numStates+1; i++) {
		for (int j = 0; j<N; j++) {
			*(vec1+j) = *(xmeas+i+j*(2*numStates+1)) - *(zbt+j);
		}
		for (int j = 0; j < numStates; j++) {
			*(vec2+j) = *(xprocess+i+j*(2*numStates+1)) - *(mubt+j);
		}
		wca = wc[i];
		dger_(&M,&N,&wca,vec2,&ione,vec1,&ione,Scc,&M);
	}

	free(vec1); free(vec2);
}

void Filter::computePredictedObservation(double *xmeas, double *sigmaPts,
		bool flag) {

	int N;

	if (flag) { // IMU
		N = imuSize;
	} else { // GPS + IMU
		N = gpsSize + imuSize;
	}

	// Temporary variables
	double *tmp1 = (double*)malloc(sizeof(double)*numStates);
	double *tmp2 = (double*)malloc(sizeof(double)*N);

	// Loop over all 2*N+1 sigma points
	for (int i = 0; i < 2*numStates+1; i++) {
		// Create local copy of sigma point vector
		for (int j = 0; j < numStates; j++) {
			*(tmp1+j) = *(sigmaPts+i+j*(2*numStates+1));
		}
		ac->getMeasModelState(tmp1,tmp2,dt,flag);
		// Store result from process model
		for (int j = 0; j < N; j++) {
			*(xmeas+i+j*(2*numStates+1)) = *(tmp2+j);
		}
	}
	free(tmp1); free(tmp2);

}

void Filter::computePredictedMeasurementUncertainty(double *xmeas,
		double *zbt, double *St, bool flag) {

	int N;
	double wca;

	if (flag) { // IMU
		N = imuSize;
	} else { // GPS + IMU
		N = gpsSize+imuSize;
	}

	double *vec1 = (double*)malloc(sizeof(double)*N);

	for (int j = 0; j < N; j++) {
		for (int k=j; k < N; k++) {
			*(St+k+j*N) = 0.0;
		}
	}

	for (int i = 0; i < 2*numStates+1; i++) {
		for (int j = 0; j<N; j++) {
			*(vec1+j) = *(xmeas+i+j*(2*numStates+1)) - *(zbt+j);
		}

		wca = wc[i];
		dger_(&N,&N,&wca,vec1,&ione,vec1,&ione,St,&N);
	}

	if (flag) { // IMU
		for (int i = 0; i < N; i++) {
			*(St + i + i*N) += Q_IMU[i]*dt;
		}
	} else { // GPS + IMU
		for (int i = 0; i < N; i++) {
			*(St + i + i*N) += Q_IMU_GPS[i]*dt;
		}
	}

	free(vec1);
}

void Filter::computePredictedMeasurement(double *xmeas, double *zbt, bool flag) {
	if (flag) { // IMU packet received
		for (int j = 0; j < imuSize; j++) {
			for(int i = 0; i < 2*numStates+1; i++) {
				*(zbt+j) += wm[i]*(*(xmeas+i+j*(2*numStates+1)));
			}
		}
	} else { // IMU + GPS packet received
		for (int j = 0; j < gpsSize+imuSize; j++) {
			for(int i = 0; i < 2*numStates+1; i++) {
				*(zbt+j) += wm[i]*(*(xmeas+i+j*(2*numStates+1)));
			}
		}
	}
}

void Filter::generateSigmaPointsMeasurement(double *sigmaPts, double *mubt,
		double *Sbart) {

	// Take square root of covariance matrix
	double *A = (double*)malloc(sizeof(double)*numStates*numStates);
	computeMatrixSquareRoot(A, Sbart, numStates);

	// First sigma point is the mean
	for (int i = 0; i < numStates; i++) {
		*(sigmaPts+i*(2*numStates+1)) = *(mubt+i);
	}

	// Write positive sigma pts
	for (int j = 0; j < numStates; j++) {
		for (int i = 1; i < numStates+1; i++) {
			*(sigmaPts+i+j*(2*numStates+1)) = *(mubt+j)+gamma*(*(Sbart+j+(i-1)*numStates));
		}
	}

	// Write negative sigma pts
	for (int j = 0; j < numStates; j++) {
		for (int i = numStates+1; i < 2*numStates+1; i++) {
			*(sigmaPts+i+j*(2*numStates+1)) = *(mubt+j)-gamma*(*(Sbart+j+(i-(numStates+1))*numStates));
		}
	}

	free(A);
}

void Filter::computePredictedVariance(double *xprocess, double *mubt, double *Sbart) {

	int N = numStates;
	double wca;
	double *vec1 = (double*)malloc(sizeof(double)*numStates);

	for (int j = 0; j < numStates; j++) {
		for (int k=j; k<numStates; k++) {
			*(Sbart+j+k*numStates) = 0.0;
		}
	}

	for (int i = 0; i < 2*numStates+1; i++) {
		for (int j = 0; j<numStates; j++) {
			*(vec1+j) = *(xprocess+i+j*(2*numStates+1)) - *(mubt+j);
		}

		wca = wc[i];
		dger_(&N,&N,&wca,vec1,&ione,vec1,&ione,Sbart,&N);
	}

	for (int i = 0; i < numStates; i++) {
		*(Sbart + i + i*numStates) += R[i]*dt;
	}

	free(vec1);
}

void Filter::computePredictedMean(double *xprocess, double *mubt) {
	for (int j = 0; j < numStates; j++) {
		for(int i = 0; i < 2*numStates+1; i++) {
			*(mubt+j) += wm[i]*(*(xprocess+i+j*(2*numStates+1)));
		}
	}
}

void Filter::propagatePtsThroughProcessModel(double *xprocess, double *sigmaPts) {

	// Temporary variables
	double *tmp1 = (double*)malloc(sizeof(double)*numStates);
	double *tmp2 = (double*)malloc(sizeof(double)*numStates);

	// Loop over all 2*N+1 sigma points
	for (int i = 0; i < 2*numStates+1; i++) {
		// Create local copy of sigma point vector
		for (int j = 0; j < numStates; j++) {
			*(tmp1+j) = *(sigmaPts+i+j*(2*numStates+1));
		}
		// Do not run N,E,D, or biases through aero solver
		ac->getProcessModelState(tmp1,tmp2,dt);
		// Store result from process model
		for (int j = 0; j < numStates; j++) {
			*(xprocess+i+j*(2*numStates+1)) = *(tmp2+j);
			// Store first column
		}
	}
	free(tmp1); free(tmp2);
}

void Filter::generateSigmaPointsProcess(double *sigmaPts, double *xcurr) {

	// Take square root of process covariance matrix
	double *A = (double*)malloc(sizeof(double)*numStates*numStates);
	double *Ptmp = (double*)malloc(sizeof(double)*numStates*numStates);

	/* initialize the matrix, upper triangular part */
	for (int i = 0; i < numStates; ++i) {
		for (int j = i; j < numStates; ++j) {
			Ptmp[j*numStates+i] = P[i][j];
		}
	}

	computeMatrixSquareRoot(A, Ptmp, numStates);

	// First sigma point is current state
	for (int i = 0; i < numStates; i++) {
		*(sigmaPts+i*(2*numStates+1)) = *(xcurr+i);
	}

	// Write positive sigma pts
	for (int j = 0; j < numStates; j++) {
		for (int i = 1; i < numStates+1; i++) {
			*(sigmaPts+i+j*(2*numStates+1)) = *(xcurr+j)+gamma*(*(A+j+(i-1)*numStates));
		}
	}

	// Write negative sigma pts
	for (int j = 0; j < numStates; j++) {
		for (int i = numStates+1; i < 2*numStates+1; i++) {
			*(sigmaPts+i+j*(2*numStates+1)) = *(xcurr+j)-gamma*(*(A+j+(i-(numStates+1))*numStates));
		}
	}
	free(A); free(Ptmp);
}

void Filter::getCurrentState(double *xcurr) {
	*xcurr      = sv.Np;
	*(xcurr+1)  = sv.Ep;
	*(xcurr+2)  = sv.Dp;
	*(xcurr+3)  = sv.VNb;
	*(xcurr+4)  = sv.VEb;
	*(xcurr+5)  = sv.VDb;
	*(xcurr+6)  = sv.q0;
	*(xcurr+7)  = sv.q1;
	*(xcurr+8)  = sv.q2;
	*(xcurr+9)  = sv.q3;
	*(xcurr+10) = sv.p;
	*(xcurr+11) = sv.q;
	*(xcurr+12) = sv.r;
	*(xcurr+13) = pv.VNi_wind;
	*(xcurr+14) = pv.VEi_wind;
	*(xcurr+15) = pv.VDi_wind;
	*(xcurr+16) = pv.abx;
	*(xcurr+17) = pv.aby;
	*(xcurr+18) = pv.abz;
	*(xcurr+19) = pv.wbx;
	*(xcurr+20) = pv.wby;
	*(xcurr+21) = pv.wbz;
	*(xcurr+22) = pv.barob;
	*(xcurr+23) = pv.diffb;
}

// Input A is the array the square root will be written to
// Input P is the square symmetric PSD matrix to take the square root of
void computeMatrixSquareRoot(double *A, double *P, int N) {

	double *B, *W, *Z, *WORK, abstol;
	int *ISUPPZ, *IWORK;
	int i, j, izero, lwork, liwork, info, M;

	izero = 0;

	/* initialize the matrix, upper triangular part */
	for (i = 0; i < N; ++i) {
		for (j = i; j < N; ++j) {
			A[j*N+i] = *(P+j*N+i);
		}
	}

	/* allocate space for the output parameters and workspace arrays */
	W = (double*)malloc(N * sizeof(double));
	Z = (double*)malloc(N * N * sizeof(double));
	ISUPPZ = (int*)malloc(2 * N * sizeof(int));
	WORK = (double*)malloc(26 * N * sizeof(double));
	IWORK = (int*)malloc(10 * N * sizeof(int));
	abstol = dlamch_("S");
	lwork = 26*N;
	liwork = 10*N;

	/* get the eigenvalues and eigenvectors */
	dsyevr_("V", "A", "U", &N, A, &N, &izero, &izero, &izero, &izero,
			&abstol, &M, W, Z, &N, ISUPPZ, WORK, &lwork, IWORK, &liwork, &info);

	/* allocate and initialise a new matrix B=Z*D */
	B = (double*)malloc(N * N * sizeof(double));
	for (j = 0; j < N; ++j) {
		double lambda = sqrt(W[j]);
		for (i = 0; i < N; ++i) {
			B[j*N+i] = Z[j*N+i]*lambda;
		}
	}

	/* calculate the square root A=B*Z^T */
	double dzero = 0.0;
	double done = 1.0;
	dgemm_("N", "T", &N, &N, &N, &done, B, &N, Z, &N, &dzero, A, &N);

	free(W); free(Z); free(ISUPPZ); free(WORK); free(IWORK);;
}

void Filter::updateState(double *mut) {
	sv.Np  = *mut;
	sv.Ep  = *(mut+1);
	sv.Dp  = *(mut+2);
	sv.VNb = *(mut+3);
	sv.VEb = *(mut+4);
	sv.VDb = *(mut+5);
	sv.q0  = *(mut+6);
	sv.q1  = *(mut+7);
	sv.q2  = *(mut+8);
	sv.q3  = *(mut+9);
	sv.p   = *(mut+10);
	sv.q   = *(mut+11);
	sv.r   = *(mut+12);
	pv.VNi_wind = *(mut+13);
	pv.VEi_wind = *(mut+14);
	pv.VDi_wind = *(mut+15);
	pv.abx = *(mut+16);
	pv.aby = *(mut+17);
	pv.abz = *(mut+18);
	pv.wbx = *(mut+19);
	pv.wby = *(mut+20);
	pv.wbz = *(mut+21);
	pv.barob = *(mut+22);
	pv.diffb = *(mut+23);
}
